import requests
import json
import time
import webbrowser

def get_acc_t(ak,sk):
    host = 'https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=%s&client_secret=%s' % (ak, sk)
    response = requests.get(host)
    if response:
        json_str = response.content.decode('utf-8')
    access_token = eval(json_str)['access_token']
    return access_token

def main(acc_t,text):
    url = f"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/txt2img?access_token={acc_t}"
    payload = json.dumps({
        "text":text,
        "resolution": "1024*1024",
        "style": "油画"  #可选择多种图片风格
    })
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    return response.text

def get_img(acc_t,task_id):
    url = f"https://aip.baidubce.com/rpc/2.0/ernievilg/v1/getImg?access_token={acc_t}&taskId={task_id}"
    payload =""
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }
    response = requests.request("POST", url, headers=headers, data=payload)
    return response.text

if __name__ == '__main__':
    ak = ""  #输入你自己的百度API Key
    sk = ""  #输入你自己的百度Secert Key
    access_token=get_acc_t(ak,sk)
    while 1:
        id=input("【1】进行AI作画\n【2】退出\n")
        if id=="1":
            text = input("请输入AI作图的文字描述（不大于100字）：\n")
            create_id=main(access_token,text)
            task_id=eval(create_id)["data"]["taskId"]
            print(task_id)
            img_result=get_img(access_token,task_id)
            while eval(img_result)["data"]["img"]=="":
                print("图片生成中，请稍后......")
                time.sleep(10)
                img_result = get_img(access_token, task_id)
            webbrowser.open(eval(img_result)["data"]["img"])
        elif id=="2":
            break
